module mMatrix

        use mPrecision

        implicit none

        integer, parameter :: MAX_DIM = 70

        type tMatrix
                real(D) :: data(MAX_DIM, MAX_DIM)
                real(D) :: inverse(MAX_DIM, MAX_DIM)
                integer :: nColumns = 0
                integer :: nRows = 0
                logical :: inverseCalculated = .false.
                logical :: verbose = .false.
        end type tMatrix

        interface new
                module procedure new_matrix
                module procedure new_matrix0
        end interface new

        interface write
                module procedure write_matrix
        end interface write

        interface invert
                module procedure invert_matrix
        end interface invert

        interface writeInverse
                module procedure writeInverse_matrix
        end interface writeInverse

        interface element
                module procedure element_matrix
        end interface element

        interface addToElement
                module procedure addToElement_matrix
        end interface addToElement

        interface getArray
                module procedure getArray_matrix
        end interface getArray

        interface getVerbose
                module procedure getVerbose_matrix
        end interface getVerbose

        interface getInverse
                module procedure getInverse_matrix
        end interface getInverse

        interface inverseElement
                module procedure inverseElement_matrix
        end interface inverseElement

        interface getInverseCalculated
                module procedure getInverseCalculated_matrix
        end interface getInverseCalculated

        interface operator (*)
                module procedure multiply_D
        end interface

        private
        public :: tMatrix
        public :: new, write
        public :: element
        public :: addToElement
        public :: getArray, getInverse, inverseElement
        public :: writeInverse, getInverseCalculated
        public :: getVerbose
        public :: deleteColumn, deleteRow
        public :: operator (*)

contains

        subroutine new_matrix(matrix, data, verbose)

                use mMisc

                type(tMatrix), intent(out) :: matrix
                real(D), intent(in) :: data(:, :)
                logical, intent(in), optional :: verbose

                call assert( size(data, 1) <= MAX_DIM, 'Matrix too large. Increase MAX_DIM.' )
                call assert( size(data, 2) <= MAX_DIM, 'Matrix too large. Increase MAX_DIM.' )

                matrix%data( 1:size(data, 1), 1:size(data, 2) ) = data
                matrix%nColumns = size(data, 1)
                matrix%nRows = size(data, 2)

                if( present(verbose) ) matrix%verbose = verbose
                
        end subroutine new_matrix

        subroutine new_matrix0(matrix, nColumns, nRows, verbose)

                use mMisc

                type(tMatrix), intent(out) :: matrix
                integer, intent(in) :: nColumns, nRows
                logical, intent(in), optional :: verbose

                call assert( nColumns <= MAX_DIM, 'Matrix too large. Increase MAX_DIM.' )
                call assert( nRows <= MAX_DIM, 'Matrix too large. Increase MAX_DIM.' )

                matrix%data( 1:nColumns, 1:nRows ) = 0.0_D
                matrix%nRows = nRows
                matrix%nColumns = nColumns

                if( present(verbose) ) matrix%verbose = verbose
                
        end subroutine new_matrix0

        real(D) function element_matrix(matrix, row, column)

                type(tMatrix), intent(in) :: matrix
                integer, intent(in) :: row, column

                element_matrix = matrix%data(row, column)

        end function element_matrix

        subroutine addToElement_matrix(matrix, column, row, data)

                type(tMatrix), intent(inout) :: matrix
                integer, intent(in) :: column, row
                real(D), intent(in) :: data

                matrix%data(column, row) = matrix%data(column, row) + data

        end subroutine addToElement_matrix

        function getArray_matrix(matrix)

                type(tMatrix), intent(in) :: matrix
                real(D) :: getArray_matrix(matrix%nColumns, matrix%nRows)

                getArray_matrix = matrix%data(1:matrix%nColumns, 1:matrix%nRows)

        end function getArray_matrix

        integer pure function getNColumns(matrix)

                type(tMatrix), intent(in) :: matrix

                getNColumns = matrix%NColumns

        end function getNColumns

        integer pure function getNRows(matrix)

                type(tMatrix), intent(in) :: matrix

                getNRows = matrix%NRows

        end function getNRows

        subroutine write_matrix(matrix)

                type(tMatrix), intent(in) :: matrix

                write(6, '(<matrix%nColumns>e10.2)') matrix%data( 1:getNColumns(matrix), 1:getNRows(matrix) )

        end subroutine write_matrix

        subroutine invert_matrix(matrix)

                type(tMatrix), intent(in) :: matrix

                write(0, *) 'ERROR: general matrix inversion not yet implemented.'

        end subroutine invert_matrix

        function getInverse_matrix(matrix)

                type(tMatrix), intent(in) :: matrix
                real(D) :: getInverse_matrix( getNColumns(matrix), getNRows(matrix) )

                if( .not.getInverseCalculated(matrix) ) call invert(matrix)

                getInverse_matrix = matrix%inverse( 1:getNColumns(matrix), 1:getNRows(matrix) )

        end function getInverse_matrix

        logical function getInverseCalculated_matrix(matrix)

                type(tMatrix), intent(in) :: matrix

                getInverseCalculated_matrix = matrix%inverseCalculated

        end function getInverseCalculated_matrix

        subroutine writeInverse_matrix(matrix)

                type(tMatrix), intent(in) :: matrix

                if(matrix%inverseCalculated) then
                        write(6, '(<matrix%nRows>e10.2)') &
                            matrix%inverse(1:matrix%nRows, 1:matrix%nColumns)
                else
                        write(6, *) 'Inverse not calculated.'
                end if

        end subroutine writeInverse_matrix

        logical function getVerbose_matrix(matrix)

                type(tMatrix) :: matrix

                getVerbose_matrix = matrix%verbose

        end function getVerbose_matrix

        subroutine deleteColumn(matrix, column)

                type(tMatrix), intent(inout) :: matrix
                integer, intent(in) :: column

                matrix%data(column:matrix%nColumns - 1, :) = matrix%data(column + 1:matrix%nColumns, :)

                matrix%nColumns = matrix%nColumns - 1
                matrix%inverseCalculated = .false.

        end subroutine deleteColumn

        subroutine deleteRow(matrix, row)

                type(tMatrix), intent(inout) :: matrix
                integer, intent(in) :: row

                matrix%data(:, row:matrix%nRows - 1) = matrix%data(:, row + 1:matrix%nRows)
                matrix%nRows = matrix%nRows - 1

                matrix%inverseCalculated = .false.

        end subroutine deleteRow

        real(D) function inverseElement_matrix(matrix, col, row)

                type(tMatrix), intent(in) :: matrix
                integer, intent(in) :: col, row

                if(.not.matrix%inverseCalculated) call invert(matrix)

                inverseElement_matrix = matrix%inverse(col, row)

        end function inverseElement_matrix

        function multiply_D(a, matrix)

                real(D), intent(in) :: a
                type(tMatrix), intent(in) :: matrix
                type(tMatrix) :: multiply_D

                multiply_D = matrix
                multiply_D%data = a * multiply_D%data

        end function multiply_D

end module mMatrix

module mSymmetricMatrix

        use mPrecision
        use mMatrix

        implicit none

        type tSymmetricMatrix
                type(tMatrix) :: matrix
        end type tSymmetricMatrix

        interface new
                module procedure new_symmetricMatrix
                module procedure new_symmetricMatrix0
        end interface new

        interface write
                module procedure write_symmetricMatrix
        end interface write

        interface element
                module procedure element_symmetricMatrix
        end interface element

        interface addToElement
                module procedure addToElement_symmetricMatrix
        end interface addToElement

        interface getArray
                module procedure getArray_symmetricMatrix
        end interface getArray

        interface getInverse
                module procedure getInverse_symmetricMatrix
        end interface getInverse

        interface invert
                module procedure invert_symmetricMatrix
        end interface invert

        interface writeInverse
                module procedure writeInverse_symmetricMatrix
        end interface writeInverse

        interface getInverseCalculated
                module procedure getInverseCalculated_symmetricMatrix
        end interface getInverseCalculated

        interface dim
                module procedure dim_symmetricMatrix
        end interface dim

        interface operator (*)
                module procedure multiply
                module procedure multiply_D
        end interface

        interface copyToLowerTriangle
                module procedure copyToLowerTriangle_data
        end interface copyToLowerTriangle

        interface copyToUpperTriangle
                module procedure copyToUpperTriangle_data
        end interface copyToUpperTriangle

        interface getVerbose
                module procedure getVerbose_symmetricMatrix
        end interface getVerbose

        interface inverseElement
                module procedure inverseElement_symmetricMatrix
        end interface inverseElement

        private
        public :: tSymmetricMatrix, new, write
        public :: element, getArray, getInverse, inverseElement, dim
        public :: invert, writeInverse
        public :: addToElement
        public :: copyToLowerTriangle
        public :: copyToUpperTriangle
        public :: operator (*)
        public :: getVerbose
        public :: deleteIndex

contains

        subroutine new_symmetricMatrix(matrix, data, verbose)

                use mMisc

                type(tSymmetricMatrix), intent(out) :: matrix
                real(D), intent(in) :: data(:, :)
                logical, intent(in), optional :: verbose

                call assert( size(data, 1) == size(data, 2), 'Symmetric matrix should be square.' )

                if( present(verbose) ) then
                        call new(matrix%matrix, data, verbose)
                else
                        call new(matrix%matrix, data)
                end if

        end subroutine new_symmetricMatrix

        subroutine new_symmetricMatrix0(matrix, dim, verbose)

                use mMisc

                type(tSymmetricMatrix), intent(out) :: matrix
                integer, intent(in) :: dim
                logical, intent(in), optional :: verbose

                if( present(verbose) ) then
                        call new(matrix%matrix, dim, dim, verbose)
                else
                        call new(matrix%matrix, dim, dim)
                end if

        end subroutine new_symmetricMatrix0

        subroutine write_symmetricMatrix(matrix)

                type(tSymmetricMatrix), intent(in) :: matrix

                call write(matrix%matrix)

        end subroutine write_symmetricMatrix

        real(D) function element_symmetricMatrix(matrix, row, column)

                type(tSymmetricMatrix), intent(in) :: matrix
                integer, intent(in) :: row, column

                element_symmetricMatrix = element(matrix%matrix, row, column)

        end function element_symmetricMatrix

        function getArray_symmetricMatrix(matrix)

                type(tSymmetricMatrix), intent(in) :: matrix
                real(D) :: getArray_symmetricMatrix(matrix%matrix%nColumns, matrix%matrix%nRows)
                
                getArray_symmetricMatrix = getArray(matrix%matrix)

        end function getArray_symmetricMatrix

        function getInverse_symmetricMatrix(matrix)

                type(tSymmetricMatrix), intent(inout) :: matrix
                real(D) :: getInverse_symmetricMatrix(matrix%matrix%nColumns, matrix%matrix%nRows)
                
                if( .not.getInverseCalculated(matrix) ) call invert(matrix)

                getInverse_symmetricMatrix = getInverse(matrix%matrix)

        end function getInverse_symmetricMatrix

        pure integer function dim_symmetricMatrix(matrix)

                type(tSymmetricMatrix), intent(in) :: matrix
                
                dim_symmetricMatrix = matrix%matrix%nColumns

        end function dim_symmetricMatrix

        subroutine invert_symmetricMatrix(symmetricMatrix)

                use LA_PRECISION, only: WP => DP
                use F77_LAPACK, only: LA_POTRF, LA_POTRI

                type(tSymmetricMatrix), intent(inout) :: symmetricMatrix

                integer :: info

                symmetricMatrix%matrix%inverse(1:symmetricMatrix%matrix%nColumns, &
                                               1:symmetricMatrix%matrix%nColumns) = &
                    symmetricMatrix%matrix%data(1:symmetricMatrix%matrix%nColumns, &
                                                1:symmetricMatrix%matrix%nColumns)

                ! Choelsky decomposition:
                call LA_POTRF( 'L', symmetricMatrix%matrix%nColumns, &
                               symmetricMatrix%matrix%inverse(1:symmetricMatrix%matrix%nColumns, &
                                                              1:symmetricMatrix%matrix%nColumns), &
                               symmetricMatrix%matrix%nColumns, info )

                if(info /= 0) then
                        write(0, *) 'WARNING: Decomposition error ', info 
                        write(0, *) symmetricMatrix%matrix%data(:, info)
                        write(0, *) '-------------------------------------------'
                        write(0, *) symmetricMatrix%matrix%data(info, :)
                        write(0, *) '-------------------------------------------'
                        call write(symmetricMatrix%matrix)
                        stop
                end if

                ! Invert matrix:
                call LA_POTRI( 'L', symmetricMatrix%matrix%nColumns, &
                               symmetricMatrix%matrix%inverse(1:symmetricMatrix%matrix%nColumns, &
                                                              1:symmetricMatrix%matrix%nColumns), &
                               symmetricMatrix%matrix%nColumns, info)
                if(info /= 0) then
                        write(0, *) 'ERROR: Inversion error ', info
                        call write(symmetricMatrix%matrix)
                        stop
                end if

                call copyToLowerTriangle(symmetricMatrix%matrix%inverse)

                symmetricMatrix%matrix%inverseCalculated = .true.

                if(symmetricMatrix%matrix%verbose) call testInverse(symmetricMatrix)

        end subroutine invert_symmetricMatrix

        subroutine testInverse(symmetricMatrix)

                use mBLAS

                type(tSymmetricMatrix), intent(in) :: symmetricMatrix

                type(tSymmetricMatrix) :: temp
                real(D) :: answer(symmetricMatrix%matrix%nColumns, symmetricMatrix%matrix%nRows)

                temp = symmetricMatrix

                call copyToLowerTriangles(temp)

                call xSYMM( 'L', 'L', temp%matrix%nColumns, temp%matrix%nRows, 1.0_D, &
                    temp%matrix%data(1:temp%matrix%nColumns, 1:temp%matrix%nRows), temp%matrix%nColumns, &
                    temp%matrix%inverse(1:temp%matrix%nRows, 1:temp%matrix%nColumns), temp%matrix%nRows, &
                    0.0_D, answer, size(answer, 1) )

                write(0, *) 'Check inverse:'
                write(0, '(<size(answer, 1)>f7.3)') answer

        end subroutine testInverse

        subroutine writeInverse_symmetricMatrix(symmetricMatrix)

                type(tSymmetricMatrix), intent(in) :: symmetricMatrix

                call writeInverse(symmetricMatrix%matrix)

        end subroutine writeInverse_symmetricMatrix

        subroutine addToElement_symmetricMatrix(matrix, column, row, data)

                type(tSymmetricMatrix), intent(inout) :: matrix
                integer, intent(in) :: column, row
                real(D), intent(in) :: data

                call addToElement(matrix%matrix, column, row, data)

        end subroutine addToElement_symmetricMatrix

        subroutine copyToLowerTriangles(symmetricMatrix)

                type(tSymmetricMatrix), intent(inout) :: symmetricMatrix

                call copyToLowerTriangle(symmetricMatrix%matrix%data)
                call copyToLowerTriangle(symmetricMatrix%matrix%inverse)

        end subroutine copyToLowerTriangles

        subroutine copyToLowerTriangle_data(data)

                use mMisc

                real(D), intent(inout) :: data(:, :)

                integer :: i

                call assert( size(data, 1) == size(data, 2), 'Not a square matrix')

                do i = 1, size(data, 1) - 1
                        data( i, i + 1:size(data, 2) ) = data( i + 1:size(data, 1), i)
                end do

        end subroutine copyToLowerTriangle_data

        subroutine copyToUpperTriangles(symmetricMatrix)

                type(tSymmetricMatrix), intent(inout) :: symmetricMatrix

                call copyToUpperTriangle(symmetricMatrix%matrix%data)
                call copyToUpperTriangle(symmetricMatrix%matrix%inverse)

        end subroutine copyToUpperTriangles

        subroutine copyToUpperTriangle_data(data)

                use mMisc

                real(D), intent(inout) :: data(:, :)

                integer :: i

                call assert( size(data, 1) == size(data, 2), 'Not a square matrix')

                do i = 1, size(data, 1) - 1
                        data( i + 1:size(data, 2), i ) = data( i, i + 1:size(data, 1) )
                end do

        end subroutine copyToUpperTriangle_data

        function multiply(matrix1, matrix2)

                use mMisc
                use mBLAS

                type(tSymmetricMatrix), intent(in) :: matrix1, matrix2
                type(tSymmetricMatrix) :: multiply

                type(tSymmetricMatrix) :: tempMatrix1, tempMatrix2
                real(D) :: answer(matrix1%matrix%nColumns, matrix2%matrix%nRows)

                call assert(matrix1%matrix%nColumns == matrix2%matrix%nColumns, &
                    'Matrices have different dimensions.')

                write(0, *) 'WARNING: multiply function untested'

                tempMatrix1 = matrix1
                tempMatrix2 = matrix2
                
                call copyToLowerTriangles(tempMatrix1)
                call copyToLowerTriangles(tempMatrix2)

                call xSYMM( 'L', 'L', tempMatrix1%matrix%nColumns, tempMatrix1%matrix%nRows, 1.0_D, &
                    tempMatrix1%matrix%data(1:tempMatrix1%matrix%nColumns, 1:tempMatrix1%matrix%nRows), &
                    tempMatrix1%matrix%nColumns, &
                    tempMatrix2%matrix%data(1:tempMatrix2%matrix%nColumns, 1:tempMatrix2%matrix%nRows), &
                    tempMatrix2%matrix%nColumns, &
                    0.0_D, answer, size(answer, 1) )

                call new(multiply, answer)

        end function multiply

        function multiply_D(a, matrix)

                real(D), intent(in) :: a
                type(tSymmetricMatrix), intent(in) :: matrix
                type(tSymmetricMatrix) :: multiply_D

                multiply_D%matrix = a * matrix%matrix

        end function multiply_D

        logical function getVerbose_symmetricMatrix(matrix)

                type(tSymmetricMatrix), intent(in) :: matrix

                getVerbose_symmetricMatrix = getVerbose(matrix%matrix)

        end function getVerbose_symmetricMatrix

        subroutine deleteIndex(matrix, index)

                type(tSymmetricMatrix), intent(inout) :: matrix
                integer, intent(in) :: index

                type(tMatrix) :: newMatrix

                newMatrix = matrix%matrix

                call deleteColumn(newMatrix, index)
                call deleteRow(newMatrix, index)

                matrix%matrix = newMatrix

        end subroutine deleteIndex

        real(D) function inverseElement_symmetricMatrix(matrix, col, row)

                type(tSymmetricMatrix), intent(inout) :: matrix
                integer, intent(in) :: col, row

                if( .not.getInverseCalculated(matrix) ) call invert(matrix)

                inverseElement_symmetricMatrix = inverseElement(matrix%matrix, col, row)

        end function inverseElement_symmetricMatrix

        logical function getInverseCalculated_symmetricMatrix(matrix)

                type(tSymmetricMatrix), intent(in) :: matrix

                getInverseCalculated_symmetricMatrix = getInverseCalculated(matrix%matrix)

        end function getInverseCalculated_symmetricMatrix

end module mSymmetricMatrix
